import random
import numpy as np
import cv2

from typing import Tuple, List

from pathlib import Path
from sklearn.preprocessing import OneHotEncoder
from sklearn.metrics import classification_report
from sklearn.model_selection import train_test_split
from tensorflow.keras import regularizers
from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import SGD
from tensorflow.keras.layers import Dense


def read_dog_cat() -> Tuple[np.ndarray, np.ndarray]:
    url = Path('D:/bee/dataset/dog_cat')
    x_data = []
    y_data = []

    def _temp(x_list: List[np.ndarray],
              y_list: List[List[str]],
              label: str,
              *,
              image_num: int = 12500) -> None:
        for num, i in enumerate(range(image_num)):
            raw_data = cv2.imread(str(url / Path(f'{label}.{i}.jpg')))
            final_data = cv2.resize(raw_data, (32, 32))
            x_list.append(final_data.flatten())
            y_list.append([label])
            print(f'\r{label}的第{num + 1}张加载成功！', end='')
        print()

    _temp(x_data, y_data, 'cat', image_num=10000)
    _temp(x_data, y_data, 'dog', image_num=10000)
    return np.array(x_data), np.array(y_data)


def main():
    data, label = read_dog_cat()
    # 独热编码
    hot_encoder = OneHotEncoder()
    label = hot_encoder.fit_transform(label).toarray()
    data = data / 255.0
    train_x, test_x, train_y, test_y = train_test_split(data, label,
                                                        test_size=0.25)
    train_x, val_x, train_y, val_y = train_test_split(train_x, train_y,
                                                      test_size=0.1)

    # 建立模型
    model = Sequential()
    model.add(Dense(256, input_shape=(3072,), activation="sigmoid"))
    model.add(Dense(128, activation="sigmoid", kernel_regularizer=regularizers.l2(0.01)))
    model.add(Dense(2, activation="softmax"))

    # 开始训练
    print('[info]:开始训练..')
    sgd = SGD(0.01)
    model.compile(loss="categorical_crossentropy",
                  optimizer=sgd,
                  metrics=["accuracy"])
    _ = model.fit(train_x, train_y, epochs=2500, batch_size=64, validation_data=(val_x, val_y))
    # 评价
    predictions = model.predict(test_x, batch_size=64)
    print(classification_report(test_y.argmax(axis=1),
                                predictions.argmax(axis=1),
                                target_names=[str(i) for i in hot_encoder.categories_[0]]))  #


if __name__ == '__main__':
    main()
    # from sklearn.preprocessing import OneHotEncoder
    # a = np.array([['cat'], ['dog'], ['dog']])
    # one = OneHotEncoder()
    # c = one.fit_transform(a).toarray()
    # # print(b)
    # print(c)
    # print([str(i) for i in one.categories_])
    # print(one.transform([['cat']]).toarray())
